'''
Created on Mar 11, 2015

@author: dewey
'''

from __future__ import division

import numpy as np
import math
from datetime import datetime

import pystrat.units as units

METADATA_NAME = "name"
METADATA_DESCRIPTION = "description"

class DataObject(object):
    '''An object to standardize the setting and retrieval of metadata for
    all classes in this file.'''
    
    def __init__(self, _id=None, metadata=None):
        self._metadata = metadata if metadata else {}
        self.modified_observer = None
        self._id = _id
    
    def set_metadata(self, key, value):
        '''Sets metadata key to value. If value is None, the key is deleted.'''
        if value is None and key in self._metadata:
            del self._metadata[key]
        elif value is not None:
            self._metadata[key] = value
    
    def metadata(self, key):
        '''Returns the value of key in metadata, or None if the key does not exist.'''
        if key in self._metadata:
            return self._metadata[key]
        else:
            return None
    
    def set_modified(self):
        if self.modified_observer:
            self.modified_observer(self)
    
    def name(self):
        if METADATA_NAME in self._metadata:
            return self._metadata[METADATA_NAME]
        else:
            return self._id
    
    def set_name(self, value):
        self.set_metadata(METADATA_NAME, value)
    
    def copy(self):
        return eval(repr(self))
    
    def __str__(self):
        return "%s (%s)" % (self._id, type(self))
        
    def __repr__(self):
        raise NotImplementedError("__repr__ must be defined for subclass of DataObject")

class XYData(DataObject):
    '''Replicate aware XY data object (replicates for attribute values)'''
    
    def __init__(self, _id=None, metadata=None, xy=None, data=None):
        super(XYData, self).__init__(metadata)
        self._data = data if data else {}
        self.xy = xy
        
    def add_data(self, key, value, sample_id=None):
        attrs = self._data
        if value is None and key in attrs:
            del attrs[key]
        elif not value is None:
            if not key in attrs:
                attrs[key] = []
            attrs[key].append((value, sample_id))
    
    def remove_data(self, key, value_or_id):
        if key in self._data:
            keydata = self._data[key]
            for index in range(len(keydata)):
                if keydata[index][0] == value_or_id or keydata[index][1] == value_or_id:
                    keydata.pop(index)
                    break
            if len(keydata) == 0:
                del self._data[key]
    
    def data_all(self, key, datadict=None, real_only=False):
        data = datadict if datadict else self._data
        if key in data:
            keydata = data[key]
            values = []
            sample_ids = []
            for value, sample_id in keydata:
                sample_ids.append(sample_id)
                if real_only and math.isnan(value):
                    continue
                values.append(value)
            return values, sample_ids
        else:
            return None
                
    
    def data_summary(self, key):
        data = self._data
        if key in data:
            values, sample_ids = self.data_all(key, data, True)
            n = len(values)
            if n > 0:
                stdev = np.std(values)
                average = np.average(values)
                return average, n, stdev, sample_ids
            else:
                return float('nan'), 0, None, sample_ids
        else:
            return None

    def __repr__(self):
        return "%s(%s, %s, %s, %s)" % (type(self).__name__, repr(self._id),
                                            repr(self._metadata), 
                                            repr(self.xy), 
                                            repr(self._data).replace('nan', "float('nan')") )

class VerticalData(DataObject):
    '''Replicate-aware vertical (time-series, depth-series) data.'''
        
    def __init__(self, _id=None, metadata=None, key_unit=None, value_unit=None, data=None):
        super(VerticalData, self).__init__(_id, metadata)
        self.key_unit = key_unit
        self.value_unit = value_unit
        self._data = data if data else {}
        
    def __repr__(self):
        return "%s(%s, %s, %s, %s, %s)" % (type(self).__name__, repr(self._id), 
                                                     repr(self._metadata), 
                                                     repr(self.key_unit), 
                                                     repr(self.value_unit), 
                                                     repr(self._data).replace("nan", "float('nan')"))
    
    def name(self, unit=True):
        '''Convenience method to return the user readable name of the proxy'''
        name_string = self.metadata(METADATA_NAME)
        if not name_string:
            name_string = self._id
        if self.value_unit:
            name_string = "%s (%s)" % (name_string, units.label(self.value_unit))
        return name_string
    
    def add_value(self, key, value, sample_id=None):
        data = self._data
        key = float(key)
        value = float(value)
        if not key in data:
            data[key] = []
        data[key].append((value, sample_id))
    
    def depths(self):
        return list(self._data.keys())
    
    def __get_data(self, key=None, datadict=None, real_only=True):
        data = datadict if datadict else self._data
        if key in data:
            values_list = data[key]
            values = []
            sample_ids = []
            for value, sample_id in values_list:
                sample_ids.append(sample_id)
                if real_only and math.isnan(value):
                    continue
                values.append(value)
            return values, sample_ids
        else:
            return None
    
    def data_all(self, key=None, datadict=None):
        data = datadict if datadict else self._data
        if key is None:
            keys = []
            sample_ids = []
            values = []
            for datakey in sorted(data.keys()):
                values_list = data[datakey]
                for value, sample_id in values_list:
                    keys.append(datakey)
                    sample_ids.append(sample_id)
                    values.append(value)
            return keys, values, sample_ids
        else:
            key = float(key)
            return self.__get_data(key, None, False)
        
    def data_summary(self, key=None, datadict=None):
        data = datadict if datadict else self._data
        if key is None:
            keys = []
            values = []
            n_vals = []
            stdev_vals = []
            sample_id_vals = []
            
            for datakey in sorted(data.keys()):
                keydata = self.data_summary(datakey, data)

                keys.append(datakey)
                values.append(keydata[0])
                n_vals.append(keydata[1])
                stdev_vals.append(keydata[2])
                string_ids = [str(item) for item in keydata[3]]
                sample_id_vals.append(", ".join(string_ids))
            
            return keys, values, n_vals, stdev_vals, sample_id_vals
                
        else:
            key = float(key)
            if key in data:
                values, sample_ids = self.__get_data(key, data, True)
                n = len(values)
                if n > 0:
                    stdev = np.std(values)
                    average = np.average(values)
                    return average, n, stdev, sample_ids
                else:
                    return float('nan'), 0, None, sample_ids
                
            else:
                return None

def _interpolate(x1, x2, y1, y2, x_value):
    ratio_between = float(x_value-x1) / float(x2 - x1)
    return y1 + ratio_between * (y2 - y1)

UNIT_14C_YEARS = "14C_Years"
UNIT_CAL_YEARS_BP = "Cal_Years_BP"
UNIT_CAL_YEARS_BEFORE_1950 = "Cal_Years_Pre_1950"
UNIT_YEAR_AD = "Year_AD"

class LinearAgeDepthModel(DataObject):
    '''An age-depth model based on a list of depths with a list of corresponding ages.
    The constructor does not check to see if the relationship is a 1-1 relationship, which
    is necessary for predictable results. Depths must be ordered when passed into the
    constructor. Currently pretty inefficient.'''
    
    def __init__(self, _id=None, metadata=None, depths=None, ages=None, age_unit=None, depth_unit=None):
        super(LinearAgeDepthModel, self).__init__(_id, metadata)
        self.__depths = list(depths) if depths else []
        self.__ages = list(ages) if ages else []
        self.age_unit = age_unit
        self.depth_unit = depth_unit
    
    def __repr__(self):
        return "LinearAgeDepthModel(%s, %s, %s, %s, %s, %s)" % (repr(self._id), 
                                                            repr(self._metadata), 
                                                            repr(self.__depths), 
                                                            repr(self.__ages), 
                                                            repr(self.age_unit),
                                                            repr(self.depth_unit))
    
    def add_agedepth(self, age, depth):
        self.__ages.append(age)
        self.__depths.append(depth)
    
    def to_vdata(self):
        vdata = VerticalData()
        ages = self.__ages
        depths = self.__depths
        for index in range(len(ages)):
            vdata.add_value(depths[index], ages[index], None)
        vdata.value_unit = self.age_unit
        vdata.key_unit = self.depth_unit
        return vdata
        
    def ages(self, list_of_depths=None):
        if list_of_depths:
            outlist = [self.__get_age(depth) for depth in list_of_depths]
            return outlist
        else:
            return list(self.__ages)
    
    def depths(self, list_of_ages=None):
        if list_of_ages:
            outlist = [self.__get_depth(age) for age in list_of_ages]
            return outlist
        else:
            return list(self.__ages)
        
    #these two should be depreceated
    def get_age(self, depth):
        return self.__get_age(depth)
    
    def get_depth(self, age):
        return self.__get_depth(age)
        
    def __get_age(self, depth):
                
        for index in range(1, len(self.__depths)):
            age1 = self.__ages[index-1]
            depth1 = self.__depths[index-1]
            age2 = self.__ages[index]
            depth2 = self.__depths[index]
            
            if depth1 <= depth <= depth2:
                return _interpolate(depth1, depth2, age1, age2, depth)
        
        min_depth = self.depth_range()[0]
        
        if depth < min_depth:
            return _interpolate(self.__depths[0], self.__depths[1], self.__ages[0], self.__ages[1], depth)
        else:
            return _interpolate(self.__depths[-2], self.__depths[-1], self.__ages[-2], self.__ages[-1], depth)
    
    def __get_depth(self, age):
        
        for index in range(1, len(self.__depths)):
            age1 = self.__ages[index-1]
            depth1 = self.__depths[index-1]
            age2 = self.__ages[index]
            depth2 = self.__depths[index]
            
            if age1 <= age <= age2:
                return _interpolate(age1, age2, depth1, depth2, age)
        
        age1, age2 = self.__ages[0], self.__ages[-1]
        age_ascending = age2 > age1
        if age_ascending:
            if age < age1:
                return _interpolate(self.__ages[0], self.__ages[1], self.__depths[0], self.__depths[1], age)
            else:
                return _interpolate(self.__ages[-2], self.__ages[-1], self.__depths[-2], self.__depths[-1], age)
        else:
            if age > age2:
                return _interpolate(self.__ages[0], self.__ages[1], self.__depths[0], self.__depths[1], age)
            else:
                return _interpolate(self.__ages[-2], self.__ages[-1], self.__depths[-2], self.__depths[-1], age)
        
    def depth_range(self):
        return min(self.__depths), max(self.__depths)
    
    def age_range(self):
        return min(self.__ages), max(self.__ages)

    def is_inverted(self):
        return self.age_unit != UNIT_YEAR_AD
    
    def auto_detect_unit(self):
        mindepth = min(self.__depths)
        maxdepth = max(self.__depths)
        shallowage = self.__ages[self.__depths.index(mindepth)]
        deepage = self.__ages[self.__depths.index(maxdepth)]
        
        slope = (deepage - shallowage) / (maxdepth - mindepth)
        if slope < 0:
            age_unit = UNIT_YEAR_AD
        else:
            if shallowage < 0:
                age_unit = UNIT_CAL_YEARS_BEFORE_1950
            else:
                age_unit = UNIT_CAL_YEARS_BP
        self.age_unit = age_unit
    
    def convert_ages(self, new_age_unit, current_year=None):
        if new_age_unit == self.age_unit():
            return LinearAgeDepthModel(self.__depths, self.__ages, self.age_unit)
        if current_year is None:
            current_year = int(datetime.now().year)
        if new_age_unit == UNIT_14C_YEARS or self.age_unit() == UNIT_14C_YEARS:
            raise ValueError("Cannot convert to or from 14C Years, external calibration is needed")
        old_unit = self.age_unit()
        
        new_ages = None
        
        if old_unit == UNIT_CAL_YEARS_BEFORE_1950:
            if new_age_unit == UNIT_YEAR_AD:
                new_ages = [years_before_1950*-1+1950 for years_before_1950 in self.__ages]
            elif new_age_unit == UNIT_CAL_YEARS_BP:
                new_ages = [years_before_present+(current_year-1950) for years_before_present in self.__ages]
        elif old_unit == UNIT_YEAR_AD:
            if new_age_unit == UNIT_CAL_YEARS_BEFORE_1950:
                new_ages = [(1950-year_ad) for year_ad in self.__ages]
            elif new_age_unit == UNIT_CAL_YEARS_BP:
                new_ages = [(current_year-year_ad) for year_ad in self.__ages]
        elif old_unit == UNIT_CAL_YEARS_BP:
            if new_age_unit == UNIT_CAL_YEARS_BEFORE_1950:
                new_ages = [cal_years_bp-(current_year-1950) for cal_years_bp in self.__ages]
            if new_age_unit == UNIT_YEAR_AD:
                new_ages = [(current_year - cal_years_bp) for cal_years_bp in self.__ages]
        else:
            raise ValueError("Unrecognized age unit: %s" % old_unit)
        return LinearAgeDepthModel(self.__depths, new_ages, new_age_unit)

class VerticalSegmentData(DataObject):
    
    def __init__(self, _id=None, metadata=None, data=None):
        super(VerticalSegmentData, self).__init__(_id, metadata)
        self._data = data if data else {}
    
    def __repr__(self):
        return "VerticalSegmentData(%s, %s, %s)" % (repr(self._id), repr(self._metadata), repr(self._data))
    
    def add_segment(self, minmax, attributes):
        true_minmax = (float(min(minmax)), float(max(minmax)))
        self._data(true_minmax)
    
    def segment_data(self, minmax=None):
        if minmax is None:
            keys = sorted(self._data.keys(), key=lambda range_tup: range_tup[0])
            attributes = []
            for key in keys:
                attributes.append(self._data[key])
        else:
            true_minmax = (float(min(minmax)), float(max(minmax)))
            if true_minmax in self._data:
                return self._data[true_minmax]
            else:
                return None
    
    def segment_at_value(self, value):
        keys = sorted(self._data.keys(), key=lambda range_tup: range_tup[0])
        for key in keys:
            if key[0] <= value < key[1]:
                return key, self._data[key]
    
    def modify_range(self, oldrange, newrange):
        true_minmax = (float(min(oldrange)), float(max(oldrange)))
        if true_minmax in self._data:
            value = self._data[true_minmax]
            del self._data[true_minmax]
            true_new_minmax = (float(min(newrange)), float(max(newrange)))
            self._data[true_new_minmax] = value
        else:
            raise ValueError("Range %s not found in ranges" % true_minmax)
        
class Dendrogram(DataObject):
    
    def __init__(self, _id, metadata, sub_trees, dissimilarity, dissimilarity_increase, num_values, central_key, key_unit=None):
        super(Dendrogram, self).__init__(_id, metadata)
        self.sub_trees = sub_trees
        self.dissimilarity = dissimilarity
        self.dissimilarity_increase = dissimilarity_increase
        self.central_key = central_key
        self.n = num_values
        self.key_unit = key_unit
        self.value_unit = None
    
    def __repr__(self):
        return "Dendrogram(%s, %s, %s, %s, %s, %s, %s, %s)" % (repr(self._id), repr(self._metadata), 
                                                            repr(self.sub_trees), repr(self.dissimilarity), 
                                                            repr(self.dissimilarity_increase), repr(self.n), 
                                                            repr(self.central_key), repr(self.key_unit))

    def __str__(self, level=0):
        spaces = "    " * level
        if self.sub_trees is None:
            return "%sDatapoint key:%s\n" % (spaces, self.central_key)
        else:
            stringout = "%sTree depth:%s dissimilarity:%s\n" % (spaces, self.central_key, self.dissimilarity)
            for tree in self.sub_trees:
                stringout += tree.__str__(level+1)
            return stringout

if __name__ == "__main__":
    #test vert data
    
    dta = VerticalData()
    dta.add_value(0.0, 1.0, "DTA0")
    dta.add_value(5, 2.0, "DTA5")
    dta.add_value(10, 3.0, "DTA10")
    dta.add_value(10, 2.0, "DTA10B")
    dta.add_value(10, 2.1, "DTA10C")
    dta.add_value(15, float('nan'), "DTA15")
    dta.add_value(20, float('nan'), "DTA20A")
    dta.add_value(20, 14.9, "DTA20B")
    dta.add_value(20, 15.1, "DTA20C")
    
    print(dta.data_all(15))
    print(dta.data_all(20))
    print(dta.data_summary(15))
    print(dta.data_summary(20))
    
    print(dta.data_all())
    print(dta.data_summary())
    print(repr(dta))
    objectcopy = eval(repr(dta))
